#!/bin/bash

BASE_PATH=$PWD
DOCKER_PATH="$BASE_PATH/docker"
ENV_FILE="$BASE_PATH/.env"
RUN_FILE_VERBOSE=1

# 颜色和图标定义
# --------------------------
COLOR_RESET="\033[0m"
COLOR_SUCCESS="\033[1;32m"  # 亮绿色
COLOR_FAILURE="\033[1;31m"  # 亮红色
COLOR_INFO="\033[1;36m"     # 青色

ICON_SUCCESS="✔"
ICON_FAILURE="✘"
ICON_INFO="ℹ️"

# --------------------------
# 封装输出函数
# --------------------------
print_success() {
  echo -e "${COLOR_SUCCESS} $1${COLOR_RESET}"
}

print_error() {
  echo -e "${COLOR_FAILURE} $1${COLOR_RESET}" >&2
}

print_info() {
  echo -e "${COLOR_INFO} $1${COLOR_RESET}"
}

print_success_with_icon() {
  echo -e "${COLOR_SUCCESS}${ICON_SUCCESS} $1${COLOR_RESET}"
}

print_error_with_icon() {
  echo -e "${COLOR_FAILURE}${ICON_FAILURE} $1${COLOR_RESET}" >&2
}

print_info_with_icon() {
  echo -e "${COLOR_INFO}${ICON_INFO} $1${COLOR_RESET}"
}

start_db() {
	if [ ! -d $DATA_HOME/clickhouse ]; then
		mkdir -p $DATA_HOME/clickhouse/data
		mkdir -p $DATA_HOME/clickhouse/sql
		mkdir -p $DATA_HOME/clickhouse/log
		mkdir -p $DATA_HOME/clickhouse/config
		cp $DOCKER_PATH/conf/clickhouse/config.xml $DATA_HOME/clickhouse/config/
		cp $DOCKER_PATH/conf/clickhouse/users.xml $DATA_HOME/clickhouse/config/
	fi

	if [ ! -d $DATA_HOME/mysql ]; then
		mkdir -p $DATA_HOME/mysql/data
		mkdir -p $DATA_HOME/mysql/log
		mkdir -p $DATA_HOME/mysql/conf
		cp $DOCKER_PATH/conf/mysql/my.cnf $DATA_HOME/mysql/conf/
	fi

	docker-compose --env-file $BASE_PATH/.env -f $DOCKER_PATH/docker-compose-db.yml up -d
}

ch_client_exec() {
	$CH_EXEC --database $CH_DATABASE --user $CH_USERNAME --password $CH_PASSWORD $@
}

ch_client_exec_root() {
	$CH_EXEC --password $CH_ROOT_PASSWORD $@
}

ch_run_query() {
	$CH_EXEC --database $CH_DATABASE --user $CH_USERNAME --password $CH_PASSWORD --query "$@"
}

ch_drop_table() {
	exists=$(ch_run_query "exists table $1")
	if [ "$exists" -eq 1 ]; then
		ch_run_query "drop table $1"
	fi
}

ch_drop_all_tables() {
	# 获取所有表名
	TABLES=$(ch_run_query "SELECT name FROM system.tables
    WHERE database = '$CH_DATABASE'
      AND name NOT LIKE '.inner%'")

	# 遍历并删除
	for table in $TABLES; do
		echo "Dropping table $table..."
		ch_drop_table $table
	done
}

ch_run_single_sql_file_root() {
	if [ ! -f $1 ]; then
		print_error "failed [file not found]"
		return
	fi
	if [ $USE_DOCKER -eq 1 ]; then
		cp $1 $DATA_HOME/clickhouse/sql
		base_file_name=$(basename "$1")
		target_file_name="/var/local/sql/$base_file_name"
	else
		target_file_name="$1"
	fi
	msg=$(ch_client_exec_root --queries-file "$target_file_name" 2>&1)
	[[ "$msg" == *"ERROR"* ]] && echo -e "\n$msg"
}

ch_run_single_sql_file() {
	if [ $RUN_FILE_VERBOSE -eq 1 ]; then
		echo -n "run $1       ..."
	fi
	if [ ! -f $1 ]; then
		print_error "failed [file not found]"
		return
	fi
	if [ $USE_DOCKER -eq 1 ]; then
		cp $1 $DATA_HOME/clickhouse/sql
		base_file_name=$(basename "$1")
		target_file_name="/var/local/sql/$base_file_name"
	else
		target_file_name="$1"
	fi
	msg=$(ch_client_exec --queries-file "$target_file_name" 2>&1)
	if [ $RUN_FILE_VERBOSE -eq 1 ]; then
		[[ "$msg" == *"ERROR"* ]] && (print_error "failed"; echo "$msg") \
			|| print_success "done"
	else
		[[ "$msg" == *"ERROR"* ]] && echo -e "\n$msg"
	fi
}

ch_run_sql_files() {
	for sql_file in $@; do		
		ch_run_single_sql_file $sql_file
	done
}

check_ch_table() {
	exists=$(ch_run_query "exists table $1") 
	[ "$exists" -eq 1 ] && print_success "done" \
		|| print_error "failed"
}

mysql_client_exec() {
	$MYSQL_EXEC -u$MYSQL_USERNAME -p$MYSQL_PASSWORD -D $MYSQL_DATABASE $@
}

mysql_client_exec_root() {
	$MYSQL_EXEC -uroot -p$MYSQL_ROOT_PASSWORD $@
}

mysql_run_query() {
	$MYSQL_EXEC -u$MYSQL_USERNAME -p$MYSQL_PASSWORD -D $MYSQL_DATABASE -BN -e "$@" 2>/dev/null
}

mysql_run_single_sql_file_root() {
	if [ ! -f $1 ]; then
		print_error "failed [file not found]"
		return
	fi
	msg=$(mysql_client_exec_root < $1 2>&1)
	[[ "$msg" == *"ERROR"* ]] && echo -e "\n$msg"
}

mysql_run_single_sql_file() {
	if [ $RUN_FILE_VERBOSE -eq 1 ]; then
		echo -n "run $1       ..."
	fi
	if [ ! -f $1 ]; then
		print_error "failed [file not found]"
		return
	fi
	msg=$(mysql_client_exec < $1 2>&1)
	if [ $RUN_FILE_VERBOSE -eq 1 ]; then
		[[ "$msg" == *"ERROR"* ]] && (print_error "failed"; echo "$msg") \
			|| print_success "done"
	else
		[[ "$msg" == *"ERROR"* ]] && echo -e "\n$msg"
	fi
}

mysql_run_sql_files() {
	for file in $@;
	do
		mysql_run_single_sql_file $file
	done
}

check_mysql_table() {
	exists=$(mysql_run_query "show tables like '$1'")
	[ -n "$exists" ] && print_success "done" \
		|| print_error "failed"
}

init_ch_db() {
	# disable verbose output
	RUN_FILE_VERBOSE=0

	echo -n "init clickhouse db  ..."
	ch_run_single_sql_file_root $BASE_PATH/sql/clickhouse/root.sql
	ch_run_single_sql_file "$BASE_PATH/sql/clickhouse/ddl.sql"
	check_ch_table "access_logs"
}

init_mysql_db() {
	# disable verbose output
	RUN_FILE_VERBOSE=0

	echo -n "init mysql db       ..."
	mysql_run_single_sql_file_root $BASE_PATH/sql/mysql/root.sql
	mysql_run_single_sql_file $BASE_PATH/sql/mysql/ddl.sql
	check_mysql_table "zms_audit_log"
}


if [ -f "$ENV_FILE" ]; then
    source $ENV_FILE
fi

if [ -z "$DATA_HOME" ]; then
    echo "DATA_HOME not set"
    exit 1
fi

if [ $# -eq 0 ]; then
	echo "Usage: $0 <start|stop|init|ch-init|ch-run|mysql-init|mysql-run>"
   	exit 1
fi

if command -v ch &>/dev/null; then
	array=(${CH_ADDR//:/ })
	host=${array[0]}
	port=${array[1]}
    CH_EXEC="ch client --host $host --port $port"
	USE_DOCKER=0
else
    CH_EXEC="docker exec -i clickhouse-server ch client"
	USE_DOCKER=1
fi

if command -v mysql &>/dev/null; then
	array=(${MYSQL_ADDR//:/ })
	host=${array[0]}
	port=${array[1]}
    MYSQL_EXEC="mysql -h $host -P $port"
else
    MYSQL_EXEC="docker exec -i mysql-server mysql"
fi

case $1 in
	"start")
		start_db
		;;
	"stop")
		docker-compose -f $DOCKER_PATH/docker-compose-db.yml down
		;;
	"init")
		init_mysql_db
		init_ch_db
		;;
	"mysql-init")
		init_mysql_db
		;;
	"mysql-run")
		shift
		mysql_run_sql_files $@
		;;
	"ch-init")
		init_ch_db
		;;
	"ch-run")
		shift
		ch_run_sql_files $@
		;;		
	"ch-drop")
		ch_drop_all_tables
		;;
	*)
		echo "unknown command $1"
		;;
esac
